import json
import os
from abc import ABCMeta, abstractmethod

import cv2
from torch.utils.data import Dataset

from ..common.types import *


class AbstractDataset(Dataset, metaclass=ABCMeta):
    """Apply template design pattern to provide an abstraction for any kind of dataset.

    """

    def __init__(self, root: str, partition: str = "train", lazy_load=False):
        self.root = root
        self.partition = partition

        if not lazy_load:
            self.load()

    @abstractmethod
    def load(self):
        pass

    @abstractmethod
    def __len__(self) -> int:
        pass

    @abstractmethod
    def __getitem__(self, index: int) -> Tuple:
        pass


class CustomClsDataset(AbstractDataset):
    """Custom dataset for classification

    See examples/dataset
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.metadata = None

    def load(self):
        metadata_filename = os.path.join(self.root, "metadata", self.partition + ".json")
        self.metadata = json.load(open(metadata_filename, 'r'))

    def __len__(self) -> int:
        return len(self.metadata)

    def __getitem__(self, index: int) -> ClsDataItem:
        filename = self.metadata[index]['filename']
        label = self.metadata[index]['label']
        image = cv2.imread(filename)

        return image, label
